#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LangChain Agent工具定义
将现有的MCP搜索、本地知识库、OPGG数据封装为LangChain Tool
"""

import asyncio
from typing import Type, Optional, Any, Dict, List
from pydantic import BaseModel, Field
from langchain.tools import BaseTool
from langchain_core.callbacks import AsyncCallbackHandler
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 导入现有功能
from mcp_web_search import SmartWebSearchClient, format_search_results
from enum import Enum


class DataSource(Enum):
    """数据源枚举"""
    LOCAL_RAG = "local_rag"
    OPGG_MCP = "opgg_mcp"
    WEB_SEARCH = "web_search"
    HYBRID = "hybrid"
    BASIC_SEARCH = "basic_search"


class WebSearchInput(BaseModel):
    """联网搜索工具输入参数"""
    query: str = Field(description="搜索查询关键词或问题")
    engine: str = Field(default="bing", description="搜索引擎，可选：bing、baidu、csdn、duckduckgo")
    max_results: int = Field(default=5, description="最大搜索结果数量")


class LocalKnowledgeInput(BaseModel):
    """本地知识库查询输入参数"""
    query: str = Field(description="查询关键词或问题")


class OPGGDataInput(BaseModel):
    """OPGG数据查询输入参数"""
    query: str = Field(description="查询内容，如：阵容、装备、英雄等")


class AgentCallbackHandler(AsyncCallbackHandler):
    """Agent执行过程的回调处理器，用于可视化思考过程"""

    def __init__(self):
        self.thoughts = []
        self.actions = []
        self.observations = []

    async def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs) -> Any:
        """链条开始时的回调"""
        if 'input' in inputs:
            print(f"🤖 Agent开始处理查询: {inputs['input']}")
        return await super().on_chain_start(serialized, inputs, **kwargs)

    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs) -> Any:
        """链条结束时的回调"""
        if 'output' in outputs:
            print(f"✅ Agent完成处理")
        return await super().on_chain_end(outputs, **kwargs)

    async def on_agent_action(self, action: Dict[str, Any], **kwargs) -> Any:
        """Agent执行动作时的回调"""
        tool = action.get('tool', 'unknown')
        tool_input = action.get('tool_input', {})
        print(f"🔧 Agent决策: 使用工具 '{tool}'")
        print(f"📝 工具参数: {tool_input}")
        self.actions.append(action)
        return await super().on_agent_action(action, **kwargs)

    async def on_agent_finish(self, finish: Dict[str, Any], **kwargs) -> Any:
        """Agent完成时的回调"""
        output = finish.get('output', '')
        print(f"🎯 Agent最终回答长度: {len(output)} 字符")
        self.observations.append(finish)
        return await super().on_agent_finish(finish, **kwargs)


class MCPWebSearchTool(BaseTool):
    """MCP联网搜索工具封装"""

    name: str = "web_search"
    description: str = """进行联网搜索获取最新信息的工具。
    适用于：最新新闻、产品评价、公司动态、技术资讯等需要实时信息的查询。
    支持多个搜索引擎：bing(默认)、baidu、csdn、duckduckgo等。
    """
    args_schema: type[WebSearchInput] = WebSearchInput

    def __init__(self):
        super().__init__()
        object.__setattr__(self, 'search_client', SmartWebSearchClient())
        object.__setattr__(self, '_initialized', False)

    async def _initialize(self):
        """异步初始化搜索客户端"""
        if not self._initialized:
            await self.search_client.initialize()
            self._initialized = True

    def _run(self, query: str, engine: str = "bing", max_results: int = 5) -> str:
        """同步调用实现"""
        # 在异步上下文中运行同步代码
        import asyncio
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(self._arun(query, engine, max_results))

    async def _arun(self, query: str, engine: str = "bing", max_results: int = 5) -> str:
        """异步调用实现"""
        try:
            await self._initialize()

            print(f"🔍 开始联网搜索: {query}")
            print(f"🌐 搜索引擎: {engine}")

            # 执行搜索
            results = await self.search_client.search(query, engine, max_results)

            if results:
                # 格式化搜索结果
                formatted_results = format_search_results(results)

                # 生成简洁的搜索结果摘要
                summary = f"联网搜索完成，找到 {len(results)} 条相关信息：\n\n"
                for i, result in enumerate(results[:3], 1):
                    summary += f"{i}. {result.title}\n   {result.snippet[:100]}...\n"

                if len(results) > 3:
                    summary += f"\n... 还有 {len(results) - 3} 条结果"

                summary += f"\n\n详细结果：\n{formatted_results}"
                return summary
            else:
                return f"联网搜索未找到关于 '{query}' 的相关信息。建议：\n1. 使用更具体的关键词\n2. 检查拼写是否正确\n3. 尝试其他搜索引擎"

        except Exception as e:
            return f"联网搜索失败: {str(e)}\n\n建议：\n1. 检查网络连接\n2. 稍后重试\n3. 尝试其他数据源"


class LocalKnowledgeTool(BaseTool):
    """本地知识库查询工具封装"""

    name: str = "local_knowledge"
    description: str = """查询本地知识库获取基础信息的工具。
    适用于：游戏基础规则、历史数据、教学指南等不需要实时性的信息。
    使用本地文档，响应快速，信息稳定。
    """
    args_schema: type[LocalKnowledgeInput] = LocalKnowledgeInput

    def __init__(self):
        super().__init__()
        object.__setattr__(self, 'basic_search', None)
        object.__setattr__(self, '_initialized', False)

    async def _initialize(self):
        """异步初始化本地知识库"""
        if not self._initialized:
            # 导入基础搜索模块
            from start import BasicTextSearch
            self.basic_search = BasicTextSearch()
            self._initialized = True

    def _run(self, query: str) -> str:
        """同步调用实现"""
        import asyncio
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(self._arun(query))

    async def _arun(self, query: str) -> str:
        """异步调用实现"""
        try:
            await self._initialize()

            print(f"📚 查询本地知识库: {query}")

            # 执行本地搜索
            results = self.basic_search.search(query)

            if results:
                formatted_results = "\n".join([f"• {result}" for result in results])
                return f"本地知识库查询结果：\n\n{formatted_results}\n\n💡 本地知识库包含基础规则和历史数据，如需最新信息建议使用联网搜索。"
            else:
                return f"本地知识库中未找到关于 '{query}' 的信息。\n\n💡 建议：\n1. 尝试使用联网搜索获取最新信息\n2. 使用更具体的关键词\n3. 检查相关概念的不同表述"

        except Exception as e:
            return f"本地知识库查询失败: {str(e)}"


class OPGGDataTool(BaseTool):
    """OPGG实时数据查询工具封装"""

    name: str = "opgg_data"
    description: str = """获取OPGG实时游戏数据和Meta信息的工具。
    适用于：最新阵容排行、装备数据、英雄强度、版本Meta分析等游戏相关查询。
    提供实时的游戏数据和统计信息。
    """
    args_schema: type[OPGGDataInput] = OPGGDataInput

    def __init__(self):
        super().__init__()
        object.__setattr__(self, 'opgg_client', None)
        object.__setattr__(self, '_initialized', False)

    async def _initialize(self):
        """异步初始化OPGG客户端"""
        if not self._initialized:
            from start import OPGGMCPClient
            self.opgg_client = OPGGMCPClient()
            await self.opgg_client.connect()
            self._initialized = True

    def _run(self, query: str) -> str:
        """同步调用实现"""
        import asyncio
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(self._arun(query))

    async def _arun(self, query: str) -> str:
        """异步调用实现"""
        try:
            await self._initialize()

            print(f"📊 查询OPGG实时数据: {query}")

            # 执行OPGG查询
            result = await self.opgg_client.query_tft_data(query)

            if result and result.content:
                return f"OPGG实时数据：\n\n{result.content}\n\n💡 OPGG数据实时更新，反映当前游戏Meta状况。"
            else:
                return f"OPGG数据查询未找到关于 '{query}' 的相关信息。\n\n💡 建议：\n1. 确认查询内容与云顶之弈相关\n2. 尝试使用联网搜索获取更广泛信息\n3. 查询本地知识库了解基础概念"

        except Exception as e:
            return f"OPGG数据查询失败: {str(e)}\n\n💡 当前使用模拟数据，建议检查网络连接。"


class SystemInfoTool(BaseTool):
    """系统信息工具"""

    name: str = "system_info"
    description: str = """获取系统信息和帮助的工具。
    当用户询问系统功能、使用方法或需要帮助时使用。
    """

    def _run(self, query: str = "") -> str:
        """系统信息查询"""
        info = """
🎯 云顶之弈智能助手 - 系统信息

🔧 可用工具：
• web_search - 联网搜索获取最新信息
• local_knowledge - 本地知识库查询
• opgg_data - OPGG实时游戏数据

💡 使用建议：
• 最新新闻/产品评价 → 使用 web_search
• 游戏基础规则 → 使用 local_knowledge
• 游戏Meta数据 → 使用 opgg_data
• 复杂查询 → 让Agent自主决策

📋 查询示例：
• "小米汽车最新评价"
• "云顶之弈基础规则"
• "当前版本什么阵容最强"
• "帮我分析最新的科技趋势"

🤖 AI Agent会根据你的查询自主选择合适的工具并生成回答！
        """
        return info.strip()


# 工具集合
TOOLS = [
    MCPWebSearchTool(),
    LocalKnowledgeTool(),
    OPGGDataTool(),
    SystemInfoTool()
]


def get_tool_by_name(tool_name: str) -> Optional[BaseTool]:
    """根据名称获取工具"""
    for tool in TOOLS:
        if tool.name == tool_name:
            return tool
    return None


def get_available_tools() -> List[Dict[str, str]]:
    """获取可用工具列表"""
    return [
        {
            "name": tool.name,
            "description": tool.description
        }
        for tool in TOOLS
    ]


# 测试代码
async def test_tools():
    """测试所有工具"""
    print("🧪 测试LangChain工具...")

    # 测试联网搜索
    web_search_tool = MCPWebSearchTool()
    print("\n1. 测试联网搜索:")
    result = await web_search_tool._arun("人工智能最新发展")
    print(result[:200] + "..." if len(result) > 200 else result)

    # 测试本地知识库
    local_tool = LocalKnowledgeTool()
    print("\n2. 测试本地知识库:")
    result = await local_tool._arun("云顶之弈规则")
    print(result[:200] + "..." if len(result) > 200 else result)

    # 测试OPGG数据
    opgg_tool = OPGGDataTool()
    print("\n3. 测试OPGG数据:")
    result = await opgg_tool._arun("当前版本阵容")
    print(result[:200] + "..." if len(result) > 200 else result)

    print("\n✅ 工具测试完成!")


if __name__ == "__main__":
    asyncio.run(test_tools())